import gradio as gr
from core.query.query_processor import QueryProcessor
from config.config import (
    API_ENDPOINT, LOCAL_MODEL_NAME, 
    REMOTE_API_KEY, REMOTE_API_ENDPOINT, REMOTE_MODEL_NAME,
    FILE_PATH, EXPLANATION_FILE
)
import pandas as pd

# 固定参数配置
file_path = FILE_PATH
explanation_file = EXPLANATION_FILE
local_api_endpoint = API_ENDPOINT
local_model_name = LOCAL_MODEL_NAME
remote_api_key = REMOTE_API_KEY
remote_api_endpoint = REMOTE_API_ENDPOINT
remote_model_name = REMOTE_MODEL_NAME

def display_output(user_input, use_remote):
    if use_remote:
        # 使用远程模型
        api_key = remote_api_key
        base_url = remote_api_endpoint
        model_name = remote_model_name
    else:
        # 使用本地模型
        api_key = None
        base_url = local_api_endpoint
        model_name = local_model_name
    
    processor = QueryProcessor(
        user_input=user_input, 
        file_path=file_path, 
        api_key=api_key, 
        base_url=base_url, 
        local_model_name=local_model_name, 
        remote_model_name=remote_model_name, 
        explanation_file=explanation_file, 
        use_remote=use_remote
    )
    
    processor.process()

    history = processor.session_manager.get_conversation_history()
    if not history:
        return "No conversation history available.", None
    
    # 获取对话历史的最后一条记录
    last_history = history[-1]

    # 根据存在的键返回相应的输出
    if "execution_result" in last_history:
        return None, last_history["execution_result"]
    elif "model_response" in last_history:
        return last_history["model_response"], None
    elif "generated_code" in last_history:
        return last_history["generated_code"], None
    else:
        return "No valid response found.", None

# Gradio UI 部分
def create_gradio_interface():
    with gr.Blocks() as demo:
        user_input = gr.Textbox(label="User Input")
        
        # 使用远程模型的选择框
        use_remote = gr.Checkbox(label="Use Remote Model", value=False)
        
        # 输出
        generated_code_output = gr.Textbox(label="Generated Code", interactive=False, visible=True)
        execution_result_output = gr.Textbox(label="Execution Result", interactive=False, visible=True)
        
        # 按钮
        submit_btn = gr.Button("Submit")
        
        # 事件处理
        submit_btn.click(display_output, 
                         inputs=[user_input, use_remote],
                         outputs=[generated_code_output, execution_result_output])
    
    return demo

if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch()



